#! /usr/bin/env runhugs +l
--
-- ASCII85.hs
-- Copyright (C) 2017 wang <wang@Yet-Another-MBP>
--
-- Distributed under terms of the MIT license.
--

module ASCII85 where
import           Data.Bits       (shiftL, (.|.))
import           Data.ByteString (ByteString, pack, unpack)
import           Data.Char       (ord)
import           Data.List.Split (chunksOf)
import           Data.Word       (Word32, Word8)

-- the result is little endian
toBaseN :: Integral a => Int -> a -> [Word8]
toBaseN n 0 = []
toBaseN n x = let (quot, rem) = quotRem x n
               in fromIntegral rem : toBaseN n quot

toBase85, toBase256 :: Integral a => a -> [Word8]
toBase85 = toBaseN 85
toBase256 = toBaseN (2 ^ 8)

strToWord8List :: String -> [Word8]
strToWord8List = map (fromIntegral . ord)


chunkToAscii :: [Word8] -> [Word8]
chunkToAscii [w]          = take (5 - 3) $ fourToFive [w, 0, 0, 0]
chunkToAscii [w, x]       = take (5 - 2) $ fourToFive [w, x, 0, 0]
chunkToAscii [w, x, y]    = take (5 - 1) $ fourToFive [w, x, y, 0]
chunkToAscii [0, 0, 0, 0] = [fromIntegral $ ord 'z']
chunkToAscii xs           = fourToFive xs

fourToFive :: [Word8] -> [Word8]
fourToFive = map (+ 33) . reverse . toBase85 . fourToWord32
    where fourToWord32 :: [Word8] -> Word32
          fourToWord32 = sum . zipWith wordShift [3, 2, 1, 0]
          wordShift n w = fromIntegral w `shiftL` (8 * n)

toAscii85 :: ByteString -> ByteString
toAscii85 = pack . surround . concat . map chunkToAscii . chunksOf 4 . unpack
    where surround s = strToWord8List "<~" ++ s ++ strToWord8List "~>"

------------------------

fiveToFour :: [Word8] -> [Word8]
fiveToFour = reverse . toBase256 . fiveToWord32
    where fiveToWord32 :: [Word8] -> Word32
          fiveToWord32 = sum . zipWith wordShift [4, 3 .. 0]
          wordShift n ascii = (fromIntegral ascii - 33) * 85 ^ n

travelAscii :: [Word8] -> [Word8]
travelAscii (122 : xs)     = 0 : 0 : 0 : 0 : travelAscii xs
travelAscii (a:b:c:d:e:xs) = fiveToFour [a,b,c,d,e] : travelAscii xs
travelAscii [w, x, y, z]   = take 3 $ fiveToFour [w, x, y, z, 117]
travelAscii [w, x, y]      = take 2 $ fiveToFour [w, x, y, 117, 117]
travelAscii [w, x]         = take 1 $ fiveToFour [w, x, 117, 117, 117]
travelAscii [w]            = error "not a chance"
travelAscii []             = []


fromAscii85 :: ByteString -> ByteString
fromAscii85 = filter unpack
